import pandas as pd
import numpy as np
import time

import peakutils
from matplotlib import pyplot as plt
from scipy.signal import find_peaks, savgol_filter

from utils.CommonFuncs import Smooth
from utils.get_iot_data import GetIotData
from function_all import FFT_spectrogram


class DiagnoseIotData:

    @staticmethod
    def diagnose_iot_signal(hz, values, motor, motorAllow, couplingType, bpfi, bpfo, bsf, fif):
        rawData = values
        fs = float(hz)
        raw_data = np.array(rawData)
        # 去直流
        data_mean = sum(raw_data) / len(raw_data)
        raw_data -= data_mean
        data = np.asarray(raw_data)
        x_axis, y_axis = FFT_spectrogram(data, fs)
        conclusion = ''
        advice = ''
        # NOTE 新增
        feature_values_dict = {}
        max_x_axis = np.max(x_axis)
        if motor != '' and 0 < float(motor) / 6 <= max_x_axis:
            motor = float(motor)
        else:
            max_xAmpl = np.max(y_axis)
            index = np.where(y_axis == max_xAmpl)
            rotate0 = x_axis[index][0] * 60
            time_range = len(data) / fs
            rotate1 = cal_rotate(time_range, data)
            if rotate0 / 6 <= min(500, max_x_axis) and rotate1 / 6 <= min(500, max_x_axis):
                motor = rotate0
            elif rotate0 / 6 > max(500, max_x_axis) and rotate1 / 6 > max(500, max_x_axis):
                motor = 980
            else:
                motor = min(rotate0, rotate1)
            # max_xAmpl = np.max(y_axis)
            # index = np.where(y_axis == max_xAmpl)[0][0]
            # motor = x_axis[index] * 60
        xampl, if_df_kong = getXampl2(x_axis, y_axis, motor, motorAllow)
        if if_df_kong == 0:
            # 不平衡
            imbalance_rate = float(xampl["1xAmpl"] / xampl.sum(axis=1))
            if couplingType != '膜片联轴器':
                feature_values_dict['一倍频特征量'] = imbalance_rate
                if imbalance_rate > 0.7:
                    conclusion += '转子不平衡'
                    advice += '调整转子平衡状态，排除轴系不平衡情况'
                else:
                    conclusion += ''
                    advice += ''
            else:
                feature_values_dict['一倍频特征量(膜片)'] = imbalance_rate
                if imbalance_rate > 0.6:
                    conclusion += '转子不平衡'
                    advice += '调整转子平衡状态，排除轴系不平衡情况'
                else:
                    conclusion += ''
                    advice += ''
            # 不对中
            centering_rate1 = float((xampl["1xAmpl"] + xampl["2xAmpl"]) / xampl.sum(axis=1))
            centering_rate2 = float(xampl["2xAmpl"] / xampl["1xAmpl"])
            feature_values_dict['一倍及二倍频特征量1'] = centering_rate1
            feature_values_dict['一倍及二倍频特征量2'] = centering_rate2
            if (centering_rate1 > 0.7 and centering_rate1 <= 0.8 and centering_rate2 <= 0.6) or (
                    centering_rate2 > 0.5 and centering_rate2 <= 0.6 and centering_rate1 <= 0.8):
                if len(conclusion) == 0:
                    conclusion += '设备对中有偏差'
                    advice += '检查联轴器状态及设备对中情况'
                else:
                    conclusion += '；设备对中有偏差'
                    advice += '；检查联轴器状态及设备对中情况'
            elif centering_rate1 > 0.8 or centering_rate2 > 0.6:
                if len(conclusion) == 0:
                    conclusion += '设备对中不良'
                    advice += '联轴器检查并重新找正'
                else:
                    conclusion += '；设备对中不良'
                    advice += '；联轴器检查并重新找正'
            else:
                conclusion += ''
                advice += ''
            # print(len(conclusion))

            # 轴承故障
            if bpfi != '' and bpfo != '' and bsf != '' and fif != '':
                return_df = pd.DataFrame(
                    columns=['name', 'ampl', 'value'])
                for i in range(1, 4):
                    tmp_dict = {}
                    tmp_dict['name'] = 'bpfi'
                    tmp_dict['ampl'] = i
                    ratio = x_axis[1] - x_axis[0]
                    energy_range_1 = int((motor / 60 * bpfi * i * (1 - motorAllow)) / ratio)
                    energy_range_2 = int((motor / 60 * bpfi * i * (1 + motorAllow)) / ratio)
                    energy = max(y_axis[energy_range_1:energy_range_2], default=0)
                    tmp_dict['value'] = round(energy, 4)
                    new_row = pd.Series(tmp_dict)
                    return_df = return_df.append(new_row, ignore_index=True)
                for i in range(1, 4):
                    tmp_dict = {}
                    tmp_dict['name'] = 'bpfo'
                    tmp_dict['ampl'] = i
                    ratio = x_axis[1] - x_axis[0]
                    energy_range_1 = int((motor / 60 * bpfo * i * (1 - motorAllow)) / ratio)
                    energy_range_2 = int((motor / 60 * bpfo * i * (1 + motorAllow)) / ratio)
                    energy = max(y_axis[energy_range_1:energy_range_2], default=0)
                    tmp_dict['value'] = round(energy, 4)
                    new_row = pd.Series(tmp_dict)
                    return_df = return_df.append(new_row, ignore_index=True)
                for i in range(1, 4):
                    tmp_dict = {}
                    tmp_dict['name'] = 'bsf'
                    tmp_dict['ampl'] = i
                    ratio = x_axis[1] - x_axis[0]
                    energy_range_1 = int((motor / 60 * bsf * i * (1 - motorAllow)) / ratio)
                    energy_range_2 = int((motor / 60 * bsf * i * (1 + motorAllow)) / ratio)
                    energy = max(y_axis[energy_range_1:energy_range_2], default=0)
                    tmp_dict['value'] = round(energy, 4)
                    new_row = pd.Series(tmp_dict)
                    return_df = return_df.append(new_row, ignore_index=True)
                for i in range(1, 4):
                    tmp_dict = {}
                    tmp_dict['name'] = 'fif'
                    tmp_dict['ampl'] = i
                    ratio = x_axis[1] - x_axis[0]
                    energy_range_1 = int((motor / 60 * fif * i * (1 - motorAllow)) / ratio)
                    energy_range_2 = int((motor / 60 * fif * i * (1 + motorAllow)) / ratio)
                    energy = max(y_axis[energy_range_1:energy_range_2], default=0)
                    tmp_dict['value'] = round(energy, 4)
                    new_row = pd.Series(tmp_dict)
                    return_df = return_df.append(new_row, ignore_index=True)
                sum_12 = return_df['value'].sum()
                bearing_rate = sum_12 / float(xampl["1xAmpl"])
                feature_values_dict['故障频率特征量'] = bearing_rate
                return_df1 = return_df[return_df['name'] == 'bpfi']
                return_df1 = return_df1.reset_index(drop=True)
                sum_1 = return_df1['value'].sum()
                return_df2 = return_df[return_df['name'] == 'bpfo']
                return_df2 = return_df2.reset_index(drop=True)
                sum_2 = return_df2['value'].sum()
                return_df3 = return_df[return_df['name'] == 'bsf']
                return_df3 = return_df3.reset_index(drop=True)
                sum_3 = return_df3['value'].sum()
                return_df4 = return_df[return_df['name'] == 'fif']
                return_df4 = return_df4.reset_index(drop=True)
                sum_4 = return_df4['value'].sum()
                if bearing_rate > 0.6 and bearing_rate <= 0.7:
                    if len(conclusion) == 0:
                        conclusion += '轴承存在早期磨损及润滑不良'
                        advice += '增加润滑，跟踪轴承振动及声音情况'
                    else:
                        conclusion += '；轴承存在早期磨损及润滑不良'
                        advice += '；增加润滑，跟踪轴承振动及声音情况'
                    if sum_1 > 0.4 * sum_12:
                        conclusion += '，并跟踪轴承内圈运行情况'
                    if sum_2 > 0.4 * sum_12:
                        conclusion += '，并跟踪轴承外圈运行情况'
                    if sum_3 > 0.4 * sum_12:
                        conclusion += '，并跟踪轴承滚动体运行情况'
                    if sum_4 > 0.4 * sum_12:
                        conclusion += '，并跟踪轴承保持架运行情况'
                elif bearing_rate > 0.7:
                    if len(conclusion) == 0:
                        conclusion += '轴承存在异常磨损'
                        advice += '增加润滑，调整轴承装配或者更换轴承'
                    else:
                        conclusion += '；轴承存在异常磨损'
                        advice += '；增加润滑，调整轴承装配或者更换轴承'
                    if sum_1 > 0.5 * sum_12:
                        conclusion += '，择机检查轴承内圈情况'
                    if sum_2 > 0.5 * sum_12:
                        conclusion += '，择机检查轴承外圈情况'
                    if sum_3 > 0.5 * sum_12:
                        conclusion += '，择机检查轴承滚动体情况'
                    if sum_4 > 0.5 * sum_12:
                        conclusion += '，择机检查轴承保持架情况'
                else:
                    conclusion += ''
                    advice += ''
            # print(len(conclusion))
            # 松动
            songdong_rate = float((xampl["3xAmpl"] + xampl["4xAmpl"] + xampl["5xAmpl"] + xampl["6xAmpl"] + xampl[
                "7xAmpl"] + xampl["8xAmpl"] + xampl["9xAmpl"] + xampl["10xAmpl"]) / xampl.sum(axis=1))
            if songdong_rate > 0.5:
                if len(conclusion) == 0:
                    conclusion += '地脚螺栓存在松动情况'
                    advice += '检查轴承座及设备地脚螺栓紧固情况'
                else:
                    conclusion += '；地脚螺栓存在松动情况'
                    advice += '；检查轴承座及设备地脚螺栓紧固情况'
            else:
                conclusion += ''
                advice += ''

        else:
            conclusion = ''
            advice = ''

        return conclusion, advice, feature_values_dict


def getXampl2(x_axis, y_axis, motor, motorAllow):
    return_df = pd.DataFrame()
    if_df_kong = 0
    arr3 = np.vstack((x_axis, y_axis))
    arr3t = arr3.T
    df = pd.DataFrame(arr3t, columns=['axisX', 'axisY'])
    bq_min = motor / 60 * (1 - motorAllow)
    bq_max = motor / 60 * (1 + motorAllow)
    df_bq1 = df[(df['axisX'] >= bq_min) & (df['axisX'] <= bq_max)]
    success = df_bq1.empty is False
    if success is True:
        df_bq1 = df_bq1.sort_values(['axisY'], ascending=False)
        df_bq1 = df_bq1.reset_index(drop=True)
        bq1 = df_bq1.loc[0, 'axisX']
        bq1_energy = df_bq1.loc[0, 'axisY']
        return_df["1xAmpl"] = [round(bq1_energy, 4)]
        for i in range(2, 11):
            bq_tmp = bq1 * i
            df_bq_tmp = df[df['axisX'] == bq_tmp]
            df_bq_tmp = df_bq_tmp.reset_index(drop=True)
            bq_tmp_energy = df_bq_tmp.loc[0, 'axisY']
            return_df[str(i) + "xAmpl"] = [round(bq_tmp_energy, 4)]
    else:
        bq_min2 = motor / 60 * (1 - motorAllow) * 2
        bq_max2 = motor / 60 * (1 + motorAllow) * 2
        df_bq2 = df[(df['axisX'] >= bq_min2) & (df['axisX'] <= bq_max2)]
        success = df_bq2.empty is False
        if success is True:
            df_bq2 = df_bq2.sort_values(['axisY'], ascending=False)
            df_bq2 = df_bq2.reset_index(drop=True)
            bq2 = df_bq2.loc[0, 'axisX']
            bq2_energy = df_bq2.loc[0, 'axisY']
            bq1 = bq2 / 2
            ratio = x_axis[1] - x_axis[0]
            df_bq1 = df[(df['axisX'] >= bq1 - ratio) & (df['axisX'] <= bq1 + ratio)]
            df_bq1 = df_bq1.sort_values(['axisY'], ascending=False)
            df_bq1 = df_bq1.reset_index(drop=True)
            bq1 = df_bq1.loc[0, 'axisX']
            bq1_energy = df_bq1.loc[0, 'axisY']
            return_df["1xAmpl"] = [round(bq1_energy, 4)]
            return_df["2xAmpl"] = [round(bq2_energy, 4)]
            for i in range(3, 11):
                bq_tmp = bq1 * i
                df_bq_tmp = df[(df['axisX'] >= bq_tmp - ratio) & (df['axisX'] <= bq_tmp + ratio)]
                df_bq_tmp = df_bq_tmp.sort_values(['axisY'], ascending=False)
                df_bq_tmp = df_bq_tmp.reset_index(drop=True)
                bq_tmp_energy = df_bq_tmp.loc[0, 'axisY']
                return_df[str(i) + "xAmpl"] = [round(bq_tmp_energy, 4)]

        else:
            if_df_kong = 1
    return return_df, if_df_kong


def cal_rotate(time_range, waveform):
    ct = len(waveform)
    t = np.linspace(0, time_range, ct)
    waveform_smooth1 = savgol_filter(waveform, window_length=31, polyorder=3)
    waveform_smooth = Smooth.smooth3(waveform_smooth1, int(ct * 0.01))
    # peaks = peakutils.indexes(waveform_smooth, thres=0.5)
    peaks, _ = find_peaks(waveform_smooth)  # 检测波峰

    # plt.figure(figsize=(15, 4))
    # plt.plot(t, waveform)
    # plt.plot(t, waveform_smooth)
    # plt.plot(t[peaks], waveform_smooth[peaks], "x", label="Peaks", color='red')
    # plt.xlabel("Time")
    # plt.ylabel("Amplitude")
    # plt.legend()
    # plt.show()

    time_diff = np.diff(t[peaks])
    time_diff_max = np.max(time_diff)
    rotate_max = 1 / time_diff_max * 60

    return rotate_max


if __name__ == '__main__':
    # log_time1 = time.time()
    # path = 'root.BGTA.BGTAIOBC00.EMFN_GWS_01.SNR.R00287.C09.V_HW_VEL'
    # now_time = '2024/01/15 14:42:41.560'
    # timeMark = '1'
    # motor = '980'
    # motorAllow = '0.02'
    # bpfo = '0.45'
    # bpfi = '0.12'
    # bsf = '0.23'
    # fif = '0.34'
    # couplingType = ''
    # if bpfo != '':
    #     bpfo = float(bpfo)
    # if bpfi != '':
    #     bpfi = float(bpfi)
    # if bsf != '':
    #     bsf = float(bsf)
    # if fif != '':
    #     fif = float(fif)
    # if motorAllow != '':
    #     motorAllow = float(motorAllow)
    # if motor != '':
    #     motor = float(motor)
    # flag = int(timeMark)
    # run_status, msg, rst = GetIotData.get_iot_signal(path, now_time, flag)
    # if run_status:
    #     hz = rst['hz']
    #     cl = rst['cl']
    #     values = rst['values']
    #     conclusion, advice = DiagnoseIotData.diagnose_iot_signal(hz, values, motor, motorAllow, couplingType, bpfi, bpfo, bsf, fif)
    #     print(f"结论：{conclusion}")
    #     print(f"建议：{advice}")
    # log_time2 = time.time()
    # print(f"用时：{log_time2 - log_time1}s")

    path_list = [
        'root.BGTA.BGTAIOBC00.EMFN_GWS_01.SNR.R00011.C01.V_VW_VEL',
        'root.BGTA.BGTAIOBC00.EMFN_GWS_01.SNR.R00004.C09.V_HW_VEL',
        'root.BGTA.BGTAIOBC00.EMFN_GWS_01.SNR.R00004.C09.V_VW_VEL',
        'root.BGTA.BGTAIOBC00.EMFN_GWS_01.SNR.R00004.C09.V_AW_VEL',
        'root.BGTA.BGTAIOBC00.EMFN_GWS_01.SNR.R00004.C10.V_HW_VEL',
        'root.BGTA.BGTAIOBC00.EMFN_GWS_01.SNR.R00004.C10.V_VW_VEL',
        'root.BGTA.BGTAIOBC00.EMFN_GWS_01.SNR.R00004.C10.V_AW_VEL',
        'root.BGTA.BGTAIOBC00.EMFN_GWS_01.SNR.R00004.C13.V_HW_VEL',
        'root.BGTA.BGTAIOBC00.EMFN_GWS_01.SNR.R00004.C13.V_VW_VEL',
        'root.BGTA.BGTAIOBC00.EMFN_GWS_01.SNR.R00004.C13.V_AW_VEL',
        'root.BGTA.BGTAIOBC00.EMFN_GWS_01.SNR.R00004.C14.V_HW_VEL',
        'root.BGTA.BGTAIOBC00.EMFN_GWS_01.SNR.R00004.C14.V_VW_VEL',
        'root.BGTA.BGTAIOBC00.EMFN_GWS_01.SNR.R00004.C14.V_AW_VEL',
        'root.BGTA.BGTAIOBC00.EMFN_GWS_01.SNR.R00014.C06.V_AW_VEL',
        'root.BGTA.BGTAIOBC00.EMFN_GWS_01.SNR.R00014.C07.V_HW_VEL',
        'root.BGTA.BGTAIOBC00.EMFN_GWS_01.SNR.R00014.C07.V_VW_VEL',
        'root.BGTA.BGTAIOBC00.EMFN_GWS_01.SNR.R00014.C07.V_AW_VEL',
        'root.BGTA.BGTAIOBC00.EMFN_GWS_01.SNR.R00014.C08.V_HW_VEL',
        'root.BGTA.BGTAIOBC00.EMFN_GWS_01.SNR.R00014.C08.V_VW_VEL',
        'root.BGTA.BGTAIOBC00.EMFN_GWS_01.SNR.R00014.C08.V_AW_VEL',
        'root.BGTA.BGTAIOBC00.EMFN_GWS_01.SNR.R00014.C09.V_HW_VEL',
        'root.BGTA.BGTAIOBC00.EMFN_GWS_01.SNR.R00014.C09.V_VW_VEL',
        'root.BGTA.BGTAIOBC00.EMFN_GWS_01.SNR.R00014.C09.V_AW_VEL',
        'root.BGTA.BGTAIOBC00.EMFN_GWS_01.SNR.R00014.C10.V_HW_VEL',
        'root.BGTA.BGTAIOBC00.EMFN_GWS_01.SNR.R00014.C10.V_VW_VEL',
        'root.BGTA.BGTAIOBC00.EMFN_GWS_01.SNR.R00014.C10.V_AW_VEL',
        'root.BGTA.BGTAIOBC00.EMFN_GWS_01.SNR.R00012.C01.V_HW_VEL',
        'root.BGTA.BGTAIOBC00.EMFN_GWS_01.SNR.R00012.C01.V_VW_VEL',
        'root.BGTA.BGTAIOBC00.EMFN_GWS_01.SNR.R00012.C01.V_AW_VEL',
        'root.BGTA.BGTAIOBC00.EMFN_GWS_01.SNR.R00012.C02.V_HW_VEL',
        'root.BGTA.BGTAIOBC00.EMFN_GWS_01.SNR.R00012.C02.V_VW_VEL',
        'root.BGTA.BGTAIOBC00.EMFN_GWS_01.SNR.R00012.C02.V_AW_VEL',
        'root.BGTA.BGTAIOBC00.EMFN_GWS_01.SNR.R00004.C17.V_HW_VEL',
        'root.BGTA.BGTAIOBC00.EMFN_GWS_01.SNR.R00004.C17.V_VW_VEL',
        'root.BGTA.BGTAIOBC00.EMFN_GWS_01.SNR.R00004.C17.V_AW_VEL',
        'root.BGTA.BGTAIOBC00.EMFN_GWS_01.SNR.R00004.C18.V_HW_VEL',
        'root.BGTA.BGTAIOBC00.EMFN_GWS_01.SNR.R00004.C18.V_VW_VEL',
        'root.BGTA.BGTAIOBC00.EMFN_GWS_01.SNR.R00004.C18.V_AW_VEL',
        'root.BGTA.BGTAIOBC00.EMFN_GWS_01.SNR.R00013.C01.V_HW_VEL',
        'root.BGTA.BGTAIOBC00.EMFN_GWS_01.SNR.R00013.C01.V_VW_VEL',
        'root.BGTA.BGTAIOBC00.EMFN_GWS_01.SNR.R00013.C01.V_AW_VEL',
        'root.BGTA.BGTAIOBC00.EMFN_GWS_01.SNR.R00013.C02.V_HW_VEL',
        'root.BGTA.BGTAIOBC00.EMFN_GWS_01.SNR.R00013.C02.V_VW_VEL',
        'root.BGTA.BGTAIOBC00.EMFN_GWS_01.SNR.R00013.C02.V_AW_VEL',
        'root.BGTA.BGTAIOBC00.EMFN_GWS_01.SNR.R00013.C05.V_HW_VEL',
        'root.BGTA.BGTAIOBC00.EMFN_GWS_01.SNR.R00013.C05.V_VW_VEL',
        'root.BGTA.BGTAIOBC00.EMFN_GWS_01.SNR.R00013.C05.V_AW_VEL',
        'root.BGTA.BGTAIOBC00.EMFN_GWS_01.SNR.R00013.C06.V_HW_VEL',
        'root.BGTA.BGTAIOBC00.EMFN_GWS_01.SNR.R00013.C06.V_VW_VEL',
        'root.BGTA.BGTAIOBC00.EMFN_GWS_01.SNR.R00013.C06.V_AW_VEL'
    ]
    now_time = '2024/01/15 14:42:41.560'
    result_df = pd.DataFrame(columns=['iotTag', 'rotate'])
    for path in path_list:
        run_status, msg, rst = GetIotData.get_iot_signal(path, now_time, 1)
        if run_status:
            rawData = rst['values']
            fs = float(rst['hz'])
            raw_data = np.array(rawData)
            data_mean = sum(raw_data) / len(raw_data)
            raw_data -= data_mean
            data = np.asarray(raw_data)
            time_range = len(data) / fs
            motor = cal_rotate(time_range, data)
            result_df = result_df.append(pd.Series([path, motor], index=result_df.columns), ignore_index=True)
    print(result_df)
